package handler

import (
	"fmt"
	"github.com/miekg/dns"
	"log/slog"
	"regexp"
	"strings"
	"sync"
	"time"
)

// DnsItem DNS记录条目
type DnsItem struct {
	hosts []string
	ttl   int64
}

func NewDnsItem() *DnsItem {
	return &DnsItem{}
}

func (d *DnsItem) Ttl() int64 {
	return d.ttl
}

func (d *DnsItem) SetTtl(ttl int64) {
	if ttl > 0 {
		d.ttl = ttl
	}
}

func (d *DnsItem) Hosts() []string {
	return d.hosts
}

func (d *DnsItem) AddHost(host string) {
	if d.hosts == nil {
		d.hosts = make([]string, 0)
	}
	d.hosts = append(d.hosts, host)
}

var mutex = &sync.Mutex{}
var dnsService *DnsService

// DnsService DNS服务对象, 提供域名解析查询
type DnsService struct {
	dnsServer      string
	dnsClient      dns.Client
	cleanCacheTime int64               // 清理dns缓存时间戳
	localCache     map[string]*DnsItem // dns缓存对象
}

// NewDnsService 构造函数, 单例模式
func NewDnsService(dnsServer string) *DnsService {
	mutex.Lock()
	defer mutex.Unlock()
	if dnsService == nil {
		if !regexp.MustCompile(`:\d{1,5}$`).MatchString(dnsServer) {
			dnsServer += ":53"
		}
		dnsService = &DnsService{dnsServer: dnsServer, dnsClient: dns.Client{}, localCache: make(map[string]*DnsItem)}
	}
	return dnsService
}

// CleanDns 清理所有DNS缓存
func (d *DnsService) CleanDns() {
	mutex.Lock()
	defer mutex.Unlock()
	d.localCache = make(map[string]*DnsItem)
}

// 清理过期DNS缓存
func (d *DnsService) cleanUpExpiredDns() {
	mutex.Lock()
	defer mutex.Unlock()
	for name, dnsItem := range d.localCache {
		if dnsItem.Ttl() < time.Now().Unix() {
			delete(d.localCache, name)
		}
	}
	d.cleanCacheTime = time.Now().Unix()
}

// 从缓存查询NDS
func (d *DnsService) queryCache(domain string) (*DnsItem, bool) {
	if time.Now().Unix() > (d.cleanCacheTime + 600) {
		// 每十分钟清理一次缓存
		go d.cleanUpExpiredDns()
	}
	if dnsItem, exists := d.localCache[domain]; exists && dnsItem.Ttl() > time.Now().Unix() {
		return dnsItem, true
	}
	return nil, false
}

// QueryForce 查询域名解析, 强制走DNS服务器
func (d *DnsService) QueryForce(domain string) (*DnsItem, error) {
	var response *dns.Msg
	var err error

	for i := 0; i < 2; i++ {
		// 尝试两次
		response, _, err = d.dnsClient.Exchange(d.buildQueryRequest(domain), d.dnsServer)
		if err == nil {
			break
		}
	}
	if err != nil {
		return nil, err
	}
	if response == nil {
		return nil, fmt.Errorf("dns response is nil")
	}

	dnsItem := NewDnsItem()
	for _, ans := range response.Answer {
		if item, ok := ans.(*dns.A); ok {
			dnsItem.AddHost(item.A.String())
			dnsItem.SetTtl(time.Now().Unix() + int64(item.Hdr.Ttl))
			slog.Info(fmt.Sprintf("Resolve domain names: %s -> %s", domain, item.A.String()))
		}
	}
	return dnsItem, nil
}

// Query 查询域名解析，先查询本地缓存，未命中就查询DNS服务器
func (d *DnsService) Query(domain string) ([]string, error) {
	if dnsItem, cache := d.queryCache(domain); cache {
		return dnsItem.Hosts(), nil
	}

	if dnsItem, err := d.QueryForce(domain); err != nil {
		return nil, err
	} else {
		d.localCache[domain] = dnsItem
		return dnsItem.Hosts(), nil
	}
}

// AutoQuery 自动解析, 参数是域名格式就查询解析, 否则返回直接返回
func (d *DnsService) AutoQuery(domain string) (string, error) {
	domainRegex := `^(?:[a-zA-Z0-9](?:[a-zA-Z0-9-]{0,61}[a-zA-Z0-9])?\.)+(?:[a-zA-Z]{2,})$`
	if regexp.MustCompile(domainRegex).MatchString(domain) {
		if hosts, err := d.Query(domain); err != nil || len(hosts) <= 0 {
			return "", err
		} else {
			return hosts[0], err
		}
	}
	return domain, nil
}

func (d *DnsService) buildQueryRequest(domain string) *dns.Msg {
	if !strings.HasSuffix(domain, ".") {
		domain = domain + "."
	}
	return new(dns.Msg).SetQuestion(domain, dns.TypeA)
}
